from grid import Grid
from scipy.spatial import ConvexHull
import numpy as np
import tetgen
# import time
# import pyvista as pv


# 计算局部梯度并筛选断层点
def get_fault_points_gradient(data, gradient_threshold):
    """
    使用梯度法识别断层点。
    """
    # 计算每个点的梯度（差分法）
    gradient = np.gradient(data, axis=0)  # 一次性计算所有维度的梯度
    # 计算总梯度（L2 范数）
    gradient_magnitude = np.linalg.norm(gradient, axis=1)
    # 筛选梯度在区间内的点
    fault_mask = (gradient_magnitude > gradient_threshold[0]) & (gradient_magnitude < gradient_threshold[1])
    fault_points = data[fault_mask]
    return fault_points


# 提取长方体的八个顶点
def get_cuboid_vertices(data):
    """
    提取长方体地质体的八个顶点。
    """
    min_x, max_x = data[:, 0].min(), data[:, 0].max()
    min_y, max_y = data[:, 1].min(), data[:, 1].max()
    min_z, max_z = data[:, 2].min(), data[:, 2].max()
    return np.array([
        [min_x, min_y, min_z], [min_x, min_y, max_z], [min_x, max_y, min_z], [min_x, max_y, max_z],
        [max_x, min_y, min_z], [max_x, min_y, max_z], [max_x, max_y, min_z], [max_x, max_y, max_z]
    ])


# 提取断层面的四个顶点
def get_fault_corners(fault_points):
    """
    使用凸包算法提取断层点的四个主要顶点。
    """
    if len(fault_points) < 4:
        raise ValueError("断层点数不足")
    # 构建凸包
    hull = ConvexHull(fault_points)
    # 从凸包顶点中提取四个角点
    hull_points = fault_points[hull.vertices]
    min_z_idx = np.argmin(hull_points[:, 2])
    max_z_idx = np.argmax(hull_points[:, 2])
    remaining_points = np.delete(hull_points, [min_z_idx, max_z_idx], axis=0)
    # 选择最小和最大z两点，以及剩余点中x方向极值点
    min_x_idx = np.argmin(remaining_points[:, 0])
    max_x_idx = np.argmax(remaining_points[:, 0])
    corners = np.array([
        hull_points[min_z_idx],  # 最低点
        hull_points[max_z_idx],  # 最高点
        remaining_points[min_x_idx],  # 最小x方向
        remaining_points[max_x_idx],  # 最大x方向
    ])
    return corners


def save_vtk_traditional(points, cells, distances, filename):
    with open(filename, 'w') as f:
        # 写入 VTK 文件头
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Tetrahedral mesh with marked points\n")
        f.write("ASCII\n")
        f.write("DATASET UNSTRUCTURED_GRID\n")
        # 写入点数据
        f.write(f"POINTS {len(points)} float\n")
        for i, point in enumerate(points):
            # 如果距离1以内，标记为 '1'，代表断层点; 否则，在坐标后加上 '0'
            if -1 <= distances[i] <= 1:
                f.write(f"{point[0]} {point[1]} {point[2]} 1\n")
            else:
                f.write(f"{point[0]} {point[1]} {point[2]} 0\n")
        # 写入单元数据
        total_entries = len(cells) * (1 + 4)  # 每个单元包含 1 个顶点数量 + 4 个顶点索引
        f.write(f"CELLS {len(cells)} {total_entries}\n")
        for cell in cells:
            f.write(f"4 {cell[0]} {cell[1]} {cell[2]} {cell[3]}\n")
        # 写入单元类型
        f.write(f"CELL_TYPES {len(cells)}\n")
        for _ in range(len(cells)):
            f.write("10\n")  # 四面体的 CELL_TYPE 是 10
    print("Mesh saved as traditional ASCII VTK format!")

def view(grid, distances):
    # 标记被切割的单元（正数在一侧，负数在另一侧，0在平面上）
    cut_cells = distances < 0
    # 分割网格为两部分
    cut_part = grid.extract_cells(cut_cells)
    remaining_part = grid.extract_cells(~cut_cells)
    # 可视化
    p = pv.Plotter()
    p.add_mesh(cut_part, show_edges=True, color="r", label="Cut Part")
    p.add_mesh(remaining_part, show_edges=True, color="lightblue", label="Remaining Part")
    p.add_legend()
    p.show()


# 主流程
def main(grid_coordinates, gradient_threshold=(5, 10), target_edge_length=15, save_filename="../Question3/Grid_E.vtk"):
    """
    主函数，提取长方体顶点和断层点并进行可视化。
    """
    # 提取长方体的八个顶点
    all_vertices = get_cuboid_vertices(grid_coordinates)
    # 筛选断层点（梯度法）
    fault_points = get_fault_points_gradient(grid_coordinates, gradient_threshold)
    # 提取断层面的四个端点
    fault_corners = get_fault_corners(fault_points)
    # 自定义面上的四个点
    p1, p2, p3 = fault_corners[0], fault_corners[1], fault_corners[3]
    # 生成凸包
    hull = ConvexHull(all_vertices)
    # 提取凸包的面的顶点索引
    all_faces = hull.simplices
    # 使用 TetGen 进行网格剖分
    tet = tetgen.TetGen(all_vertices, all_faces)
    # 修复表面使其成为流形
    tet.make_manifold(verbose=False)
    # 设置剖分密度
    max_tetrahedron_volume = (target_edge_length ** 3) / 6  # 正四面体近似体积
    tet.tetrahedralize(order=1, mindihedral=20, minratio=1.1, switches=f"-q1.2a{max_tetrahedron_volume}")
    grid = tet.grid
    # 计算法向量
    plane_normal = np.cross(p2 - p1, p3 - p1)  # 叉积计算法向量
    plane_normal = plane_normal / np.linalg.norm(plane_normal)  # 归一化法向量
    # 计算每个网格单元的中心点
    cell_centers = grid.cell_centers()
    cell_points = cell_centers.points
    # 计算点到平面的距离
    distances = np.dot(cell_points - p1, plane_normal)
    # view(grid, distances)
    # 保存为传统格式 VTK 文件
    save_vtk_traditional(grid.points, grid.cells_dict[10], distances, save_filename)


if __name__ == '__main__':
    # start = time.time()
    # 读取网格文件获取节点坐标
    grid = Grid("../Question3/Grid_E.out")
    grid_coordinates = grid.convert_corner_to_grid_coordinates()
    grid_coordinates = grid_coordinates.reshape((-1, 3))
    # 识别断层，进行四面体剖分，按照题目要求保存vtk模型文件并标注用于模拟断层的节点
    main(grid_coordinates, gradient_threshold=(5, 10), target_edge_length=15, save_filename="../Question3/Grid_E.vtk")
    # print(time.time() - start, " sec = ", (time.time() - start) / 60, "min")
